import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import gym
from matplotlib import pyplot as plt
import os


os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE'
env = gym.make('Pendulum-v1').unwrapped
'''Pendulum环境状态特征是三个，杆子的sin(角度)、cos（角度）、角速度，（状态是无限多个，因为连续），动作值是力矩，限定在[-2,2]之间的任意的小数，所以是连续的（动作也是无限个）'''
state_number = env.observation_space.shape[0]  # 3状态特征
action_number = env.action_space.shape[0]  # 1

max_action = env.action_space.high[0]
min_action = env.action_space.low[0]
torch.manual_seed(0)
env.seed(0)
RENDER = True
EP_MAX = 2500
EP_LEN = 200
GAMMA = 0.9


A_LR = 0.0001
C_LR = 0.0003
# 改200
BATCH = 128
A_UPDATE_STEPS = 10
C_UPDATE_STEPS = 10
METHOD = [
    dict(name='kl_pen', kl_target=0.01, lam=0.5),  # KL penalty
    dict(name='clip', epsilon=0.2),  # Clipped surrogate objective, find this is better
][1]  # choose the method for optimization
TrainFlag = 1

'''由于PPO也是基于A-C框架，所以我把PPO的编写分为两部分，PPO的第一部分 Actor'''
'''PPO的第一步  编写A-C框架的网络，先编写actor部分的actor网络，actor的网络有新与老两个网络'''

# optimize
# reward 标准化  0
# not use ratio clip 1
# 学习率衰减 2
# 更新AC网络的梯度裁剪 3
# 使用提前截止算法 4
use_flag = [1, 0, 1, 1, 1]
if use_flag[2]:
    AC_Clip = 0.3
    A_LR = A_LR*2
    C_LR = C_LR*2
else:
    AC_Clip = 0.4
# optimize.png是啥都没采用的情况
#



class ActorNet(nn.Module):
    def __init__(self, inp, outp):
        super(ActorNet, self).__init__()
        self.in_to_y1 = nn.Linear(inp, 100)
        self.in_to_y1.weight.data.normal_(0, 0.1)
        self.out = nn.Linear(100, outp)
        self.out.weight.data.normal_(0, 0.1)
        self.std_out = nn.Linear(100, outp)
        self.std_out.weight.data.normal_(0, 0.1)

    '''生成均值与标准差，PPO必须这样做，一定要生成分布（所以需要mean与std），不然后续学习策略里的公式写不了，DDPG是可以不用生成概率分布的'''

    def forward(self, inputstate):
        inputstate = self.in_to_y1(inputstate)
        # if not use_flag[6]:
        inputstate = F.relu(inputstate)  # 激活函数0
        # else:
        #     inputstate = torch.tanh(inputstate)  # 激活函数1
        mean = max_action * torch.tanh(self.out(inputstate))  # 输出概率分布的均值mean
        std = F.softplus(self.std_out(inputstate))  # softplus激活函数的值域>0，标准差大于0
        return mean, std


'''再编写critic部分的critic网络，PPO的critic部分与AC算法的critic部分是一样，PPO不一样的地方只在actor部分'''


class CriticNet(nn.Module):
    def __init__(self, input, output):  # 3纬数据输入 一维数据输出
        super(CriticNet, self).__init__()
        self.in_to_y1 = nn.Linear(input, 100)
        self.in_to_y1.weight.data.normal_(0, 0.1)
        self.out = nn.Linear(100, output)
        self.out.weight.data.normal_(0, 0.1)

    def forward(self, inputstate):
        # 前向计算 下面的地方调用的是这个函数
        # 输入第一层
        inputstate = self.in_to_y1(inputstate)
        # 激活函数
        # if not use_flag[6]:
        inputstate = F.relu(inputstate)  # 激活函数0
        # else:
        #     inputstate = torch.tanh(inputstate)  # 激活函数1
        # 结果输出层
        V = self.out(inputstate)
        return V


class Actor():
    def __init__(self):
        # 两个网络 用老的网络去生成一系列a
        self.old_pi, self.new_pi = ActorNet(state_number, action_number), ActorNet(state_number,
                                                                                   action_number)  # 这只是均值mean
        # 优化的目的是新actor策略网络的参数，lr为学习率，更大前期快，更小好收敛
        self.optimizer = torch.optim.Adam(self.new_pi.parameters(), lr=A_LR, eps=1e-5)
        self.cur_episode = 1

    '''第二步 编写根据状态选择动作的函数'''

    def choose_action(self, s):
        inputstate = torch.FloatTensor(s)
        mean, std = self.old_pi(inputstate)  # 输入当前的状态到老网络中，输出均值方差
        dist = torch.distributions.Normal(mean, std)
        action = dist.sample()
        # 动作值是力矩，限定在[-2,2]
        action = torch.clamp(action, min_action, max_action)
        action_logprob = dist.log_prob(action)
        # 第二次返回值是选择这个action的概率的prob
        return action.detach().numpy(), action_logprob.detach().numpy()

    '''第四步  actor网络有两个策略（更新old策略）————————把new策略的参数赋给old策略'''

    def update_oldpi(self):
        self.old_pi.load_state_dict(self.new_pi.state_dict())

    '''第六步 编写actor网络的学习函数，采用PPO，即OpenAI推出的clip形式公式'''

    def learn(self, bs, ba, adv, bap):
        bs = torch.FloatTensor(bs)  # state数组
        ba = torch.FloatTensor(ba)  # action数组
        adv = torch.FloatTensor(adv)  # td数组
        bap = torch.FloatTensor(bap)  # action的概率对数
        # loss 主要为了让adv减少 又通过ratio这个权重值让新老的策略权值，让这两个策略如果相差太大导致action半杆子打不着的话
        for _ in range(A_UPDATE_STEPS):  # 多次更新actor
            mean, std = self.new_pi(bs)  # self.new_pi每一次都变化 所以虽然用的bs一样，但是出来的分布越来越好了
            dist_new = torch.distributions.Normal(mean, std)  # 构造分布函数
            action_new_logprob = dist_new.log_prob(ba)  # 选择这个action的概率对数数组，保证两个数组的action都一致，但概率不一致
            ratio = torch.exp(action_new_logprob - bap.detach())  # 策略相除sample权重值，但是我们转化为e的x次幂变为相减
            surr1 = ratio * adv
            surr2 = torch.clamp(ratio, 1 - METHOD['epsilon'], 1 + METHOD['epsilon']) * adv
            if use_flag[1]:
                loss = -torch.min(surr1)
            else:
                loss = -torch.min(surr1, surr2)  # 上面原公式是old/new 反了所以这个取负
            loss = loss.mean()
            loss.backward()
            if use_flag[3]:  # 此参数影响很大的 AC_Clip
                nn.utils.clip_grad_norm_(self.new_pi.parameters(), AC_Clip)
            self.optimizer.step()
        # 减少学习率
        global A_LR
        if use_flag[2] and A_LR > 0.0001:  # 在1/4学习率达到下限
            A_LR = (0.0004 / 3.0) * (1 - self.cur_episode / EP_MAX)
            # A_LR = 0.0002 * (1 - self.cur_episode / EP_MAX)


class Critic():
    def __init__(self):
        self.critic_v = CriticNet(state_number, 1)  # 3，1
        self.optimizer = torch.optim.Adam(self.critic_v.parameters(), lr=C_LR, eps=1e-5)
        self.lossfunc = nn.MSELoss()
        self.cur_episode = 1

    '''第三步  编写评定动作价值的函数'''

    def get_v(self, s):
        # s是从1状态到最新状态的一个数组
        inputstate = torch.FloatTensor(s)
        # 返回的也是数组，调用forward
        return self.critic_v(inputstate)

    '''第五步  计算优势——————advantage，后面发现第五步计算出来的adv可以与第七步合为一体，所以这里的代码注释了，但是，计算优势依然算是可以单独拎出来的一个步骤'''
    # def get_adv(self,bs,br):
    #     reality_v=torch.FloatTensor(br)
    #     v=self.get_v(bs)
    #     adv=(reality_v-v).detach()
    #     return adv
    '''第七步  编写actor-critic的critic部分的learn函数，td-error的计算代码（V现实减去V估计就是td-error）'''

    def learn(self, bs, br):
        # bs为state数组，包括了这一批次所有的state
        # br为现实的价值数组，数组的每一个元素是 第i个状态点到最后累积的折扣化过的r和 比如第一个元素就是所有的价值和，第二个元素是从第二个状态到最后的价值和
        bs = torch.FloatTensor(bs)
        reality_v = torch.FloatTensor(br)
        for _ in range(C_UPDATE_STEPS):
            v = self.get_v(bs)
            td_e = self.lossfunc(reality_v, v)
            # 根据pytorch中backward（）函数的计算，当网络参量进行反馈时，梯度是累积计算而不是被替换，
            # 但在处理每一个batch时并不需要与其他batch的梯度混合起来累积计算，
            # 因此需要对每个batch调用一遍zero_grad（）将参数梯度置0.
            self.optimizer.zero_grad()
            td_e.backward()
            if use_flag[3]:
                nn.utils.clip_grad_norm_(self.critic_v.parameters(), 0.4)
            self.optimizer.step()
        # 最后得到的v是结果几次训练critic网络后得到的估计v
        global C_LR
        if use_flag[2] and C_LR > 0.0003:
            # C_LR = 0.0006 * (1 - self.cur_episode / EP_MAX)
            C_LR = 0.0004 * (1 - self.cur_episode / EP_MAX)
        return (reality_v - v).detach()


if TrainFlag == 1:
    print('PPO训练中...')
    actor = Actor()
    critic = Critic()
    all_ep_r = []
    reward_totle_episode = 0
    for episode in range(EP_MAX):  # 多轮训练
        actor.cur_episode = episode
        critic.cur_episode = episode
        print(A_LR)
        print(C_LR)
        observation = env.reset()  # 环境重置
        buffer_s, buffer_a, buffer_r, buffer_a_logp = [], [], [], []
        reward_totle = 0
        for timestep in range(EP_LEN):  # 200，我们假设200步内能完成这个实验
            if RENDER:
                env.render()
            action, action_logprob = actor.choose_action(observation)
            observation_, reward, done, info = env.step(action)
            # 这一步的数据保存起来
            buffer_s.append(observation)
            buffer_a.append(action)
            if use_flag[0]:
                buffer_r.append((reward + 8) / 8)  # normalize reward, find to be useful
            else:
                buffer_r.append((reward))  #
            buffer_a_logp.append(action_logprob)
            # 更新状态
            observation = observation_
            reward_totle += reward
            # normalize奖励
            if use_flag[0]:
                reward = (reward - reward.mean()) / (reward.std() + 1e-5)
            # PPO 更新，当做了走了几步后去更新一下，更新后保存在buffer数组的数据就清空了
            if (timestep + 1) % BATCH == 0 or timestep == EP_LEN - 1:
                # 未来的状态输入后得到未来的v
                v_observation_ = critic.get_v(observation_)
                # 数组的每一个元素是 第i个状态点到最后累积的折扣化过的r和 比如第一个元素就是所有的价值和，第二个元素是从第二个状态到最后的价值和
                discounted_r = []
                # 向前更新所有的v值
                for reward in buffer_r[::-1]:  # 把buffer_r数组逆序，因为最新的reward在最后面
                    # v_是当前状态到最后的价值积累，要去算所有状态点到最后到价值积累。反向传播到味道
                    v_observation_ = reward + GAMMA * v_observation_
                    discounted_r.append(v_observation_.detach().numpy())
                discounted_r.reverse()
                # 把数据挤压变高从1xn到nx1
                bs, ba, br, bap = np.vstack(buffer_s), np.vstack(buffer_a), np.array(discounted_r), np.vstack(
                    buffer_a_logp)
                buffer_s, buffer_a, buffer_r, buffer_a_logp = [], [], [], []
                advantage = critic.learn(bs, br)  # critic部分更新
                actor.learn(bs, ba, advantage, bap)  # actor部分更新
                actor.update_oldpi()  # pi-new的参数赋给pi-old
                # critic.learn(bs,br)
        if episode == 0:
            all_ep_r.append(reward_totle)
            reward_totle_episode = reward_totle_episode + reward_totle
        else:
            all_ep_r.append(all_ep_r[-1] * 0.9 + reward_totle * 0.1)
            reward_totle_episode = reward_totle_episode + all_ep_r[-1] * 0.9 + reward_totle * 0.1
        print("\rEp: {} |rewards: {}".format(episode, reward_totle))
        # 保存神经网络参数，50轮看一下reward的平均值 如果发现不错那么我们就结束训练

        if episode % 50 == 0 and episode > 100:  # 保存神经网络参数，是不是需要考虑一下保存reward_totle比较小的参数
            print("reward_totle_episode")
            print(reward_totle_episode / 50)
            save_data = {'net': actor.old_pi.state_dict(), 'opt': actor.optimizer.state_dict(), 'i': episode}
            torch.save(save_data, "PPO_actor.pth")
            save_data = {'net': critic.critic_v.state_dict(), 'opt': critic.optimizer.state_dict(), 'i': episode}
            torch.save(save_data, "PPO_critic.pth")
            if reward_totle_episode / 50 > -450 and use_flag[4]:  # 如果说近50轮的结果加起来平均小于200 不用训练了
                break
            reward_totle_episode = 0
    env.close()
    plt.plot(np.arange(len(all_ep_r)), all_ep_r)
    plt.xlabel('Episode')
    plt.ylabel('Moving averaged episode reward')
    use_flag_str = ''
    for i in range(0, len(use_flag)):
        if use_flag[i]:
            use_flag_str = use_flag_str + '-' + str(i + 1)
    plt.savefig("optimize{}.png".format(use_flag_str))
    plt.show()

else:
    print('PPO测试中...')
    aa = Actor()
    cc = Critic()
    checkpoint_aa = torch.load("PPO_actor.pth")
    aa.old_pi.load_state_dict(checkpoint_aa['net'])
    checkpoint_cc = torch.load("PPO_critic.pth")
    cc.critic_v.load_state_dict(checkpoint_cc['net'])
    for j in range(10):
        state = env.reset()
        total_rewards = 0
        for timestep in range(EP_LEN):  # 200
            env.render()
            action, action_logprob = aa.choose_action(state)
            new_state, reward, done, info = env.step(action)  # 执行动作
            total_rewards += reward
            state = new_state
        print("Score：", total_rewards)
    env.close()
